#!/usr/bin/env python3
"""
向量搜索工具
用于在已向量化的MD文档中搜索相关内容
"""

import json
import sys
import argparse
from pathlib import Path

# 添加当前目录到Python路径
sys.path.append(str(Path(__file__).parent))

from vectorize_md_to_milvus import MDVectorizer
import logging

def load_config(config_path: str = "config.json") -> dict:
    """加载配置文件"""
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"配置文件 {config_path} 不存在")
        return None
    except json.JSONDecodeError as e:
        print(f"配置文件格式错误: {e}")
        return None

def search_documents(query: str, top_k: int = 10, config_path: str = "config.json"):
    """搜索相关文档"""
    # 加载配置
    config = load_config(config_path)
    if not config:
        return []
    
    # 设置简单日志
    logging.basicConfig(level=logging.WARNING)
    
    try:
        # 初始化向量化处理器
        vectorizer = MDVectorizer(
            embedding_model=config["embedding"]["model"],
            embedding_url=config["embedding"]["url"],
            milvus_host=config["milvus"]["host"],
            milvus_port=config["milvus"]["port"],
            collection_name=config["milvus"]["collection_name"]
        )
        
        # 执行搜索
        results = vectorizer.search_similar(query, top_k)
        return results
        
    except Exception as e:
        print(f"搜索失败: {e}")
        return []

def format_results(results: list, show_text_length: int = 200):
    """格式化搜索结果"""
    if not results:
        print("没有找到相关结果")
        return
    
    print(f"找到 {len(results)} 个相关结果:\n")
    
    for i, result in enumerate(results, 1):
        score = result.get('score', 0)
        file_name = result.get('file_name', 'Unknown')
        chunk_index = result.get('chunk_index', 0)
        text = result.get('text', '')
        metadata = result.get('metadata', {})
        
        print(f"{i}. 相似度: {score:.4f}")
        print(f"   文件: {file_name}")
        print(f"   块索引: {chunk_index}")
        
        if metadata.get('content_type'):
            print(f"   内容类型: {metadata['content_type']}")
        
        if metadata.get('first_header'):
            print(f"   主要标题: {metadata['first_header']}")
        
        # 显示文本片段
        display_text = text[:show_text_length]
        if len(text) > show_text_length:
            display_text += "..."
        
        print(f"   内容: {display_text}")
        print("-" * 80)

def interactive_search():
    """交互式搜索模式"""
    print("向量搜索工具 - 交互模式")
    print("输入 'quit' 或 'exit' 退出")
    print("=" * 50)
    
    while True:
        try:
            query = input("\n请输入搜索查询: ").strip()
            
            if query.lower() in ['quit', 'exit']:
                print("再见!")
                break
            
            if not query:
                print("请输入有效的查询内容")
                continue
            
            # 询问返回结果数量
            try:
                top_k = input("返回结果数量 (默认5): ").strip()
                top_k = int(top_k) if top_k else 5
                top_k = max(1, min(50, top_k))  # 限制在1-50之间
            except ValueError:
                top_k = 5
            
            print(f"\n搜索: '{query}' (返回前{top_k}个结果)")
            print("=" * 50)
            
            # 执行搜索
            results = search_documents(query, top_k)
            
            # 显示结果
            format_results(results)
            
        except KeyboardInterrupt:
            print("\n\n再见!")
            break
        except Exception as e:
            print(f"搜索过程中出错: {e}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="向量搜索工具")
    parser.add_argument("query", nargs="?", help="搜索查询")
    parser.add_argument("-k", "--top-k", type=int, default=5, help="返回结果数量 (默认5)")
    parser.add_argument("-i", "--interactive", action="store_true", help="交互模式")
    parser.add_argument("-c", "--config", default="config.json", help="配置文件路径")
    parser.add_argument("--text-length", type=int, default=200, help="显示文本长度 (默认200)")
    
    args = parser.parse_args()
    
    if args.interactive:
        interactive_search()
    elif args.query:
        print(f"搜索: '{args.query}'")
        print("=" * 50)
        
        results = search_documents(args.query, args.top_k, args.config)
        format_results(results, args.text_length)
    else:
        print("用法:")
        print("  python search_documents.py '搜索关键词'")
        print("  python search_documents.py -i  # 交互模式")
        print("  python search_documents.py --help  # 显示帮助")
        return 1
    
    return 0

if __name__ == "__main__":
    exit(main())
